/*
 * Copyright OpenSearch Contributors
 * SPDX-License-Identifier: Apache-2.0
 */

package org.opensearch.sql.sql.parser;

import static java.util.Collections.emptyList;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.opensearch.sql.ast.dsl.AstDSL.aggregate;
import static org.opensearch.sql.ast.dsl.AstDSL.alias;
import static org.opensearch.sql.ast.dsl.AstDSL.distinctAggregate;
import static org.opensearch.sql.ast.dsl.AstDSL.function;
import static org.opensearch.sql.ast.dsl.AstDSL.intLiteral;
import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;

import java.util.Arrays;
import java.util.List;
import java.util.function.Function;
import org.antlr.v4.runtime.CommonTokenStream;
import org.hamcrest.FeatureMatcher;
import org.hamcrest.Matcher;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.DisplayNameGeneration;
import org.junit.jupiter.api.DisplayNameGenerator;
import org.junit.jupiter.api.Test;
import org.opensearch.sql.ast.expression.UnresolvedExpression;
import org.opensearch.sql.ast.tree.Aggregation;
import org.opensearch.sql.ast.tree.UnresolvedPlan;
import org.opensearch.sql.common.antlr.CaseInsensitiveCharStream;
import org.opensearch.sql.common.antlr.SyntaxAnalysisErrorListener;
import org.opensearch.sql.exception.SemanticCheckException;
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLLexer;
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser;
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.QuerySpecificationContext;
import org.opensearch.sql.sql.parser.context.QuerySpecification;

@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
class AstAggregationBuilderTest {

  @Test
  void can_build_group_by_clause() {
    assertThat(
        buildAggregation("SELECT state, AVG(age) FROM test GROUP BY state"),
        allOf(
            hasGroupByItems(alias("state", qualifiedName("state"))),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));
  }

  @Test
  void can_build_group_by_clause_with_scalar_expression() {
    assertThat(
        buildAggregation("SELECT ABS(age + 1) FROM test GROUP BY ABS(age + 1)"),
        allOf(
            hasGroupByItems(
                alias(
                    "ABS(+(age, 1))",
                    function("ABS", function("+", qualifiedName("age"), intLiteral(1))))),
            hasAggregators()));
  }

  @Test
  void can_build_group_by_clause_with_complicated_aggregators() {
    assertThat(
        buildAggregation("SELECT state, ABS(2 * AVG(age)) FROM test GROUP BY state"),
        allOf(
            hasGroupByItems(alias("state", qualifiedName("state"))),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));
  }

  @Test
  void can_build_group_by_clause_without_aggregators() {
    assertThat(
        buildAggregation("SELECT state FROM test GROUP BY state"),
        allOf(hasGroupByItems(alias("state", qualifiedName("state"))), hasAggregators()));
  }

  @Test
  void can_build_implicit_group_by_clause() {
    assertThat(
        buildAggregation("SELECT AVG(age), SUM(balance) FROM test"),
        allOf(
            hasGroupByItems(),
            hasAggregators(
                alias("AVG(age)", aggregate("AVG", qualifiedName("age"))),
                alias("SUM(balance)", aggregate("SUM", qualifiedName("balance"))))));
  }

  @Test
  void can_build_implicit_group_by_for_aggregator_in_having_clause() {
    assertThat(
        buildAggregation("SELECT true FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT PI() FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT ABS(1.5) FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT ABS(ABS(1.5)) FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT INTERVAL 1 DAY FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT CAST(1 AS LONG) FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));

    assertThat(
        buildAggregation("SELECT CASE WHEN true THEN 1 ELSE 2 END FROM test HAVING AVG(age) > 30"),
        allOf(
            hasGroupByItems(),
            hasAggregators(alias("AVG(age)", aggregate("AVG", qualifiedName("age"))))));
  }

  @Test
  void can_build_distinct_aggregator() {
    assertThat(
        buildAggregation("SELECT COUNT(DISTINCT name) FROM test group by age"),
        allOf(
            hasGroupByItems(alias("age", qualifiedName("age"))),
            hasAggregators(
                alias("COUNT(DISTINCT name)", distinctAggregate("COUNT", qualifiedName("name"))))));
  }

  @Test
  void should_build_nothing_if_no_group_by_and_no_aggregators_in_select() {
    assertNull(buildAggregation("SELECT name FROM test"));
  }

  @Test
  void should_replace_group_by_alias_by_expression_in_select_clause() {
    assertThat(
        buildAggregation("SELECT state AS s, name FROM test GROUP BY s, name"),
        hasGroupByItems(
            alias("state", qualifiedName("state")), alias("name", qualifiedName("name"))));

    assertThat(
        buildAggregation("SELECT ABS(age) AS a FROM test GROUP BY a"),
        hasGroupByItems(alias("ABS(age)", function("ABS", qualifiedName("age")))));
  }

  @Test
  void should_replace_group_by_ordinal_by_expression_in_select_clause() {
    assertThat(
        buildAggregation("SELECT state AS s FROM test GROUP BY 1"),
        hasGroupByItems(alias("state", qualifiedName("state"))));

    assertThat(
        buildAggregation("SELECT name, ABS(age) AS a FROM test GROUP BY name, 2"),
        hasGroupByItems(
            alias("name", qualifiedName("name")),
            alias("ABS(age)", function("ABS", qualifiedName("age")))));
  }

  @Test
  void should_report_error_for_non_integer_ordinal_in_group_by() {
    SemanticCheckException error =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT state AS s FROM test GROUP BY 1.5"));
    assertEquals("Non-integer constant [1.5] found in ordinal", error.getMessage());
  }

  @Disabled(
      "This validation is supposed to be in analyzing phase. This test should be enabled "
          + "once https://github.com/opensearch-project/sql/issues/910 has been resolved")
  @Test
  void should_report_error_for_mismatch_between_select_and_group_by_items() {
    SemanticCheckException error1 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT name FROM test GROUP BY state"));
    assertEquals(
        "Expression [name] that contains non-aggregated column is not present in group by clause",
        error1.getMessage());

    SemanticCheckException error2 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT ABS(name + 1) FROM test GROUP BY name"));
    assertEquals(
        "Expression [Function(funcName=ABS, funcArgs=[Function(funcName=+, "
            + "funcArgs=[name, Literal(value=1, type=INTEGER)])])] that contains "
            + "non-aggregated column is not present in group by clause",
        error2.getMessage());
  }

  @Test
  void should_report_error_for_non_aggregated_item_in_select_if_no_group_by() {
    SemanticCheckException error1 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT age, AVG(balance) FROM tests"));
    assertEquals(
        "Explicit GROUP BY clause is required because expression [age] "
            + "contains non-aggregated column",
        error1.getMessage());

    SemanticCheckException error2 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT ABS(age + 1), AVG(balance) FROM tests"));
    assertEquals(
        "Explicit GROUP BY clause is required because expression [ABS(+(age, 1))] "
            + "contains non-aggregated column",
        error2.getMessage());
  }

  @Test
  void should_report_error_for_group_by_ordinal_out_of_bound_of_select_list() {
    SemanticCheckException error1 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT age, AVG(balance) FROM tests GROUP BY 0"));
    assertEquals("Ordinal [0] is out of bound of select item list", error1.getMessage());

    SemanticCheckException error2 =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT age, AVG(balance) FROM tests GROUP BY 3"));
    assertEquals("Ordinal [3] is out of bound of select item list", error2.getMessage());
  }

  @Test
  void should_report_error_for_non_aggregated_item_in_select_if_only_having() {
    SemanticCheckException error =
        assertThrows(
            SemanticCheckException.class,
            () -> buildAggregation("SELECT age FROM tests HAVING AVG(balance) > 30"));
    assertEquals(
        "Explicit GROUP BY clause is required because expression [age] "
            + "contains non-aggregated column",
        error.getMessage());
  }

  private Matcher<UnresolvedPlan> hasGroupByItems(UnresolvedExpression... exprs) {
    return featureValueOf("groupByItems", Aggregation::getGroupExprList, exprs);
  }

  private Matcher<UnresolvedPlan> hasAggregators(UnresolvedExpression... exprs) {
    return featureValueOf("aggregators", Aggregation::getAggExprList, exprs);
  }

  private Matcher<UnresolvedPlan> featureValueOf(
      String name,
      Function<Aggregation, List<UnresolvedExpression>> getter,
      UnresolvedExpression... exprs) {
    Matcher<List<UnresolvedExpression>> subMatcher =
        (exprs.length == 0) ? equalTo(emptyList()) : equalTo(Arrays.asList(exprs));
    return new FeatureMatcher<UnresolvedPlan, List<UnresolvedExpression>>(subMatcher, name, "") {
      @Override
      protected List<UnresolvedExpression> featureValueOf(UnresolvedPlan agg) {
        return getter.apply((Aggregation) agg);
      }
    };
  }

  private UnresolvedPlan buildAggregation(String sql) {
    QuerySpecificationContext query = parse(sql);
    QuerySpecification querySpec = collect(query, sql);
    AstAggregationBuilder aggBuilder = new AstAggregationBuilder(querySpec);
    return aggBuilder.visit(query.fromClause().groupByClause());
  }

  private QuerySpecification collect(QuerySpecificationContext query, String sql) {
    QuerySpecification querySpec = new QuerySpecification();
    querySpec.collect(query, sql);
    return querySpec;
  }

  private QuerySpecificationContext parse(String query) {
    OpenSearchSQLLexer lexer = new OpenSearchSQLLexer(new CaseInsensitiveCharStream(query));
    OpenSearchSQLParser parser = new OpenSearchSQLParser(new CommonTokenStream(lexer));
    parser.addErrorListener(new SyntaxAnalysisErrorListener());
    return parser.querySpecification();
  }
}
